import tensorflow as tf
import keras
import numpy as np


def atrous_spatial_pyramid_pooling_keras(inputs, output_stride=8, depth=256):
    """
    atrous spatial pyramid pooling implementation with keras
    :param inputs:
    :param output_stride:
    :param batch_norm_decay:
    :param is_training:
    :param depth:
    :return:
    """
    atrous_rates = [6, 12, 18]
    if output_stride == 8:
        atrous_rates = [2 * item for item in atrous_rates]
    with tf.variable_scope('atrous_pyramid_pooling'):
        conv_1x1 = keras.layers.Conv2D(depth, (1, 1), strides=1, padding='same', activation=None)(inputs)
        conv_1x1 = keras.layers.BatchNormalization()(conv_1x1)
        conv_1x1 = keras.layers.Activation('relu')(conv_1x1)
        conv_3x3_list = []
        for item in atrous_rates:
            conv_3x3 = keras.layers.Conv2D(depth, (3, 3), strides=1, dilation_rate=item, padding='same',
                                           activation=None)(inputs)
            conv_3x3 = keras.layers.BatchNormalization()(conv_3x3)
            conv_3x3 = keras.layers.Activation('relu')(conv_3x3)
            conv_3x3_list.append(conv_3x3)
        # with tf.variable_scope("image_level_features"):
        #     # global average pooling
        #     image_level_features = keras.layers.Lambda(reduce_mean)(inputs)
        #     image_level_features = keras.layers.Conv2D(depth, (1, 1), strides=1, padding='same', activation=None)(image_level_features)
        #     image_level_features = keras.layers.BatchNormalization()(image_level_features)
        #     image_level_features = keras.layers.Activation('relu')(image_level_features)
        #     # bilinearly upsample features
        #     image_level_features = Bilinear([image_level_features, inputs])
        with tf.variable_scope('pyramid_concat'):
            net = keras.layers.Concatenate(axis=3, name='concat')(
                [conv_1x1] + conv_3x3_list)  # +[image_level_features])
            net = keras.layers.Conv2D(depth, (1, 1), strides=1, padding='same', activation=None)(net)
            net = keras.layers.BatchNormalization()(net)
            net = keras.layers.Activation('relu')(net)
            return net


def bilinear(input_tensor):
    return tf.image.resize_bilinear(input_tensor, (512, 512))


def segmentation_branch(x, num_classes):
    """
    segmentation branch
    :param x:
    :return: a tensor
    """
    with tf.variable_scope('segmentation_branch'):
        x = atrous_spatial_pyramid_pooling_keras(x, 8, 32)
        # extract output tensor of block4
        with tf.variable_scope("upsampling_logits"):
            x = keras.layers.Conv2D(num_classes, (1, 1), strides=1, padding='same', activation=None)(x)
            x = keras.layers.Lambda(bilinear)(x)
            x = keras.layers.Softmax(axis=-1, name='segmentation')(x)
    return x